import argparse
import os
import numpy as np
import torch

parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, default='LSTM', choices=['LSTM', 'CNN_LSTM', 'Seq2Seq'])
parser.add_argument('--output', type=str, default='x', choices=['x', 'y'])
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--val_rate', type=float, default=0.2)
parser.add_argument('--test_rate', type=float, default=0.1)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--step_size', type=int, default=1)
parser.add_argument('--gamma', type=float, default=0.95)
parser.add_argument('--hidden_size', type=int, default=256)
parser.add_argument('--num_layers', type=int, default=1)
args = parser.parse_args()

model_name = args.model  # LSTM CNN_LSTM Seq2Seq
result_path = 'result'
model_path = 'models'
input_size = 3  # 输入变量数
interval = 120  # 输入步长
pred_size = 20  # 预测步长
max_epochs = args.epochs  # 迭代数
val_rate = args.val_rate  # 验证集占比
test_rate = args.test_rate  # 测试集占比
hidden_size = args.hidden_size  # 隐藏层维度
num_layers = args.num_layers  # 循环层数
batch_size = args.batch_size  # 数据压缩量
learn_rate = args.lr  # 学习率
step_size = args.step_size  # 学习率递变的步长
gamma = args.gamma  # 学习率递增系数，也即每个epoch学习率变为原来的0.95
device = "cuda" if torch.cuda.is_available() else "cpu"  # GPU加速运算
output = args.output

output_n = 0 if output == 'x' else 1
output_dict = {0: 'wind_x', 1: 'wind_y'}
os.makedirs(result_path, exist_ok=True)
npfile = np.load('analysis/data.npz', allow_pickle=True)
data = npfile['d_data']
tfit = npfile['tfit']
t_date = npfile['t_date']
r_name = f'{model_name}-{output_dict[output_n]}'
test_num = int(int(len(data) * (1 - test_rate)))
train_data = data[:int(len(data) * (1 - val_rate - test_rate))]
val_data = data[int(len(data) * (1 - val_rate - test_rate)):test_num]
test_data = data[test_num:]
